// Copyright (c) MatrixFramework. All rights reserved.
// Licensed under the MIT License. See LICENSE in the project root for license information.

using System.Collections.Concurrent;
using System.Collections.Generic;
using System.Linq;
using Matrix.ServiceCommunication.Services;

namespace Matrix.ServiceCommunication.Discovery.LoadBalancing;

/// <summary>
/// 加权轮询负载均衡器
/// </summary>
public class WeightedRoundRobinLoadBalancer : ILoadBalancer
{
    private readonly ConcurrentDictionary<string, WeightCounter> _counters = new();

    public string StrategyName => "WeightedRoundRobin";

    public ServiceInstance Select(IList<ServiceInstance> instances)
    {
        if (instances == null || !instances.Any())
        {
            throw new ArgumentException("实例列表不能为空", nameof(instances));
        }

        var healthyInstances = instances.Where(i => i.Healthy && i.Enabled).ToList();
        if (!healthyInstances.Any())
        {
            throw new InvalidOperationException("没有可用的健康实例");
        }

        var serviceName = healthyInstances.First().ServiceName;
        var counter = _counters.GetOrAdd(serviceName, _ => new WeightCounter());

        lock (counter)
        {
            // 计算总权重
            var totalWeight = healthyInstances.Sum(i => Math.Max(1, i.Weight));
            var currentWeight = (counter.CurrentWeight + 1) % totalWeight;

            // 根据权重选择实例
            var accumulatedWeight = 0;
            foreach (var instance in healthyInstances)
            {
                accumulatedWeight += Math.Max(1, instance.Weight);
                if (currentWeight < accumulatedWeight)
                {
                    counter.CurrentWeight = currentWeight;
                    return instance;
                }
            }

            // 如果没有找到，返回第一个实例
            counter.CurrentWeight = currentWeight;
            return healthyInstances.First();
        }
    }

    /// <summary>
    /// 权重计数器
    /// </summary>
    private class WeightCounter
    {
        public int CurrentWeight { get; set; }
        public object Lock { get; } = new object();
    }
}